import os
import argparse

import cv2
import numpy as np
from tqdm import tqdm
import sys

from cyw_devkit import tranform_pt_pad

sys.path.append('../')  # 确保这个路径是正确的
from cyw_devkit.core.dataset import CywDataset, transform_bboxes
from cyw_devkit.core.visualizer import view3d
from cyw_devkit.core.visualizer import just_look_array, draw_pts2img, Visualizer
from tictoc import TicToc
import multiprocessing


def main(args):
    # 统计耗时
    cost = TicToc("数据可视化")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    # for idx in range(process_size):
    #     main_worker(files[0], args)

    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(dataset_path, args):
    # cyw = CywDataset(dataset_path=dataset_path)
    cyw = CywDataset(dataset_path=dataset_path, radar_dim=args.radar_dim)
    # output_file = os.path.join(args.out_path,os.path.basename(dataset_path)+'.mp4')
    # fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 设置视频编码格式（这里使用 MP4V 编码）
    # out = cv2.VideoWriter(output_file, fourcc, 30, (1080, 720))

    init_stamp = 0 # us
    boxes_all = []
    pc_all = []
    # vis = Visualizer(frame_data['lidar'].data, center_mode=False)
    # vis = Visualizer(points=None, center_mode=False)
    for idx, frame_data in enumerate(tqdm(cyw)):
        if idx==0:
            init_stamp = frame_data['lidar'].header.stamp
        #     vis = Visualizer(frame_data['lidar'].data, center_mode=False)

        # if idx<20 or idx%5!=0:
        #     continue
        # if idx>80:
        #     break
        if idx in cyw.keyframe_idxs[:120]: #判断关键帧
            now_stamp = frame_data['lidar'].header.stamp
            # now2init = cyw.data_dict['super_transform'].get_tf_from_time(source_time=now_stamp,target_time=init_stamp)
            now2init = cyw.data_dict['super_transform'].localizer.get_tf(now_stamp)


            from cyw_devkit.core.zmath import isRotationMatrix
            isRotationMatrix(now2init[:3,:3])

            pc_now = frame_data['lidar'].data[::5]
            pc_now = tranform_pt_pad(pc_now, [now2init])
            pc_all.append(pc_now)


            boxes_now = frame_data['label'].data

            # boxes_all.append(boxes_now.tensors)

            # vis = Visualizer(frame_data['lidar'].data[:,:3], center_mode=False)
            # vis = Visualizer(pc_now[:,:3], center_mode=False)
            # vis.add_bboxes(bbox3d=boxes_now.tensors)
            # vis.show()
            boxes_local = transform_bboxes(boxes_now, [now2init])
            for id in boxes_local.obj_ids:
                # if id == '1089':
                    idx = boxes_local.obj_ids.index(id)
                    boxes_all.append(boxes_local.tensors[[idx],:])

    boxes_all = np.concatenate(boxes_all, axis=0)
    pc_all = np.concatenate(pc_all, axis=0)

    view3d(pc_all[::10], boxes_all, mode='xyz')

    # vis.add_bboxes(bbox3d=boxes_all)
    # vis.show()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/home/adt/datasets",
                        help='your data root')
    parser.add_argument('--radar-dim', type=int, default=5,
                        help='cyw 5;carla 7')
    parser.add_argument('--out-path', default="/home/adt/datasets/motion/dataview",
                        help='')
    parser.add_argument('--sync-time', action='store_true',
                        help='')
    args = parser.parse_args()

    main(args)
